import os
from app.pdf_extract_kit.utils.data_preprocess import load_pdf, load_ppt
from app.project.doc_to_recommendation.utils.log_utils import get_logger


def load_file(input_data, **kwargs):
    images = []
    base_name = os.path.splitext(os.path.basename(input_data))[0]
    if (not os.path.isfile(input_data)) or (not os.path.exists(input_data)):
        raise ValueError("Input PDF file or PPT files not found. Check the path: {}".format(input_data))
    # Determine the type of input data and process accordingly
    if input_data.lower().endswith(('.pdf')):
        # If input is a single image file
        images = load_pdf(input_data)
    elif input_data.lower().endswith(('.ppt', '.pptx')):
        images = load_ppt(input_data, **kwargs)
    else:
        raise ValueError("Unsupported input data format: {}".format(input_data))
    return images, base_name

class BaseTask:

    @property
    def task_name(self):
        return "BaseTask"

    def __init__(self, model):
        self.model = model

    def load_images(self, input_data):
        """
        Loads images from a single image path or a directory containing multiple images.

        Args:
            input_data (str): Path to a single image file or a directory containing image files.

        Returns:
            list: List of paths to all images to be predicted.
        """
        images = []

        if os.path.isdir(input_data):
            # If input_data is a directory, check for nested directories
            for root, dirs, files in os.walk(input_data):
                if dirs:
                    raise ValueError("Input directory should not contain nested directories: {}".format(input_data))
                for file in files:
                    if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                        image_path = os.path.join(root, file)
                        images.append(image_path)
                images = sorted(images)
                break  # Only process the top-level directory
        else:
            # Determine the type of input data and process accordingly
            if input_data.lower().endswith(('.png', '.jpg', '.jpeg')):
                # If input is a single image file
                images = [input_data]
            else:
                raise ValueError("Unsupported input data format: {}".format(input_data))

        return images

    def process_flow(self, images: list):
        pass

    def visualize_images(self, images, results, basename, save_dir):
        pass

    def load_pdf_images(self, input_data):
        """
        Loads images from a single PDF file or directory containing multiple PDF files.

        Args:
            input_data (str): Path to a single PDF file or a directory containing PDF files.

        Returns:
            dict: Dictionary with image IDs (formed by PDF path and page number) as keys and corresponding PIL.Image objects as values.
                  Note: Loading multiple PDFs at once is not recommended due to high memory consumption. Consider processing one PDF at a time externally using loops or multithreading.
        """

        return load_file(input_data)

class TaskFlow:
    def __init__(self, task_list: list['BaseTask'], config: dict):
        self.task_list = task_list
        self.config = config
        self.log = get_logger("TaskFlow")

    def process(self, input_path, save_dir=None):
        images, base_name = load_file(input_path, **self.config.get("ppt_load", {}))
        log = self.log
        results = {}
        for task in self.task_list:
            log.info(f"================开始执行{task.task_name}任务================")
            result = task.process_flow(images)
            results.update({task.task_name : result})
            log.info(f"================任务{task.task_name}执行完成================")
            if self.config.get(task.task_name, {}).get('visualize', False) and save_dir:
                log.info(f"================开始{task.task_name}任务可视化================")
                task.visualize_images(images, result, base_name, save_dir)
                log.info(f"================任务{task.task_name}可视化完成================")
        return results, images